import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 解决pyplot绘图时标题中文乱码
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 制造数据
x_data = np.linspace(0, 6.283, 2000)[:, np.newaxis]
y_data = np.sin(x_data)


def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)


def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)


# 占位
x = tf.placeholder(tf.float32, [None, 1], name='x')
y = tf.placeholder(tf.float32, [None, 1], name='y')
keep_prob = tf.placeholder(tf.float32, name='keep_prob')

# first fully layer
w_fc1 = weight_variable([1, 2048])
b_fc1 = bias_variable([2048])
h_fc1 = tf.nn.relu(tf.nn.bias_add(tf.matmul(x, w_fc1), b_fc1))
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

# second fully layer
w_fc2 = weight_variable([2048, 2048])
b_fc2 = bias_variable([2048])
h_fc2 = tf.nn.relu(tf.nn.bias_add(tf.matmul(h_fc1_drop, w_fc2), b_fc2))
h_fc2_drop = tf.nn.dropout(h_fc2, keep_prob)

# output layer
w_fc3 = weight_variable([2048, 1])
b_fc3 = bias_variable([1])
y_output = tf.nn.bias_add(tf.matmul(h_fc2_drop, w_fc3), b_fc3, name='y_output')

# 定义损失函数和优化器
loss = tf.reduce_mean(tf.square(y_output - y_data))
train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)

real = tf.reshape(y, [-1, 1])
correct_prediction = tf.equal(y_output, real)
correct_prediction = tf.cast(correct_prediction, tf.float32)
accuracy = tf.reduce_mean(correct_prediction, name='accuracy')

with tf.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    step = 1
    train_loss = 1
    while train_loss > 0.0001:
        sess.run(train_step, feed_dict={x: x_data, y: y_data, keep_prob: 0.75})
        if step % 100 == 0:
            pred, acc, train_loss = sess.run([y_output, accuracy, loss],
                                             feed_dict={x: x_data, y: y_data, keep_prob: 1.00})
            print('step:%d,train loss:%f,accuracy:%f' % (step, train_loss, acc))
            plt.plot(x_data, y_data, label='sin(x)', color='r')
            plt.plot(x_data, pred, label='BP神经网络拟合', color='b')
            plt.title("使用BP神经网络拟合正弦函数（第{}步）".format(step))
            plt.legend(loc=1)
            plt.xlabel("x")
            plt.ylabel("y")
            plt.savefig('result/第{}步.png'.format(step))
            plt.clf()
        step += 1
